#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PaddleOCR 快速识别工具
OCR模型保持常驻状态，快速处理图片
"""

import os
import time
from pathlib import Path

from paddleocr import PaddleOCR

# 全局OCR对象
ocr = None

def init_ocr():
    """初始化OCR模型"""
    global ocr
    if ocr is None:
        print("正在初始化PaddleOCR模型...")
        ocr = PaddleOCR(
            use_doc_orientation_classify=False,
            use_doc_unwarping=False,
            use_textline_orientation=False
        )
        print("PaddleOCR模型初始化完成！")
    return ocr

def quick_recognize():
    """快速识别images文件夹中的所有图片"""
    print("=== 快速OCR识别 ===\n")
    
    # 获取OCR实例
    ocr = init_ocr()
    
    # 设置路径
    input_dir = "images"
    output_dir = "results"
    
    # 创建输出目录
    Path(output_dir).mkdir(exist_ok=True)
    
    # 支持的图片格式
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.webp']
    
    # 获取所有图片文件
    input_path = Path(input_dir)
    if not input_path.exists():
        print(f"错误：图片文件夹不存在: {input_dir}")
        return
    
    image_files = []
    for file_path in input_path.iterdir():
        if file_path.is_file() and file_path.suffix.lower() in image_extensions:
            image_files.append(file_path)
    
    if not image_files:
        print(f"在目录 {input_dir} 中未找到支持的图片文件")
        return
    
    print(f"找到 {len(image_files)} 个图片文件，开始识别...\n")
    
    # 处理每个图片
    success_count = 0
    total_start_time = time.time()
    
    for i, image_file in enumerate(image_files, 1):
        print(f"处理: {i}/{len(image_files)} - {image_file.name}", end=" ")
        
        try:
            # 记录开始时间
            start_time = time.time()
            
            # 执行OCR识别
            result = ocr.predict(str(image_file))
            
            # 记录结束时间
            end_time = time.time()
            processing_time = end_time - start_time
            
            # 提取识别结果
            if result and len(result) > 0:
                res = result[0]
                
                # 从res字段中提取数据
                try:
                    json_data = res.json
                    if 'res' in json_data:
                        res_data = json_data['res']
                        if isinstance(res_data, dict):
                            rec_texts = res_data.get('rec_texts', [])
                            rec_scores = res_data.get('rec_scores', [])
                            
                            if rec_texts:
                                # 计算平均置信度
                                avg_confidence = sum(rec_scores) / len(rec_scores) if rec_scores else 0.0
                                
                                # 合并所有文本
                                combined_text = '\n'.join(rec_texts)
                                
                                # 生成输出文件名
                                base_name = image_file.stem
                                txt_filename = base_name + '_ocr_result.txt'
                                img_filename = base_name + '_ocr_result.jpg'
                                json_filename = base_name + '_ocr_result.json'
                                
                                # 保存txt文件
                                txt_file_path = Path(output_dir) / txt_filename
                                with open(txt_file_path, 'w', encoding='utf-8') as f:
                                    f.write(f"图片文件: {image_file.name}\n")
                                    f.write(f"识别时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
                                    f.write(f"处理耗时: {processing_time:.2f} 秒\n")
                                    f.write(f"平均置信度: {avg_confidence:.4f}\n")
                                    f.write(f"识别文本数量: {len(rec_texts)}\n")
                                    f.write("-" * 50 + "\n")
                                    f.write(combined_text)
                                
                                # 保存img和json文件（使用官方方法）
                                try:
                                    res.save_to_img(str(output_dir))
                                    res.save_to_json(str(output_dir))
                                    print(f"✓ ({processing_time:.1f}s, 置信度: {avg_confidence:.3f}) - 已保存txt/img/json")
                                except Exception as save_error:
                                    print(f"✓ ({processing_time:.1f}s, 置信度: {avg_confidence:.3f}) - 已保存txt，img/json保存失败")
                                
                                success_count += 1
                            else:
                                print("✗ 未识别到文字")
                        else:
                            print("✗ 数据格式错误")
                    else:
                        print("✗ 未找到结果数据")
                        
                except Exception as e:
                    print(f"✗ 提取失败: {str(e)}")
            else:
                print("✗ 未获得识别结果")
                
        except Exception as e:
            print(f"✗ 处理失败: {str(e)}")
    
    total_time = time.time() - total_start_time
    print(f"\n处理完成！成功: {success_count}/{len(image_files)} 个文件")
    print(f"总耗时: {total_time:.2f} 秒，结果已保存到 {output_dir} 文件夹")

if __name__ == "__main__":
    quick_recognize()
